﻿#include "stdafx.h"
#include <stdio.h>
#include <windows.h>
#include <tlhelp32.h>
#include "EVSuperProcess.h"

#define NT_SUCCESS(Status) ((NTSTATUS)(Status) >= 0)
#define STATUS_SUCCESS ((NTSTATUS)0x00000000)
#define SystemProcessesAndThreadsInformation 5
#define NTAPI __stdcall



// 线程状态的枚举常量
typedef enum _THREAD_STATE
{
    StateInitialized,
    StateReady, // 准备状态
    StateRunning, // 运行状态
    StateStandby, // 
    StateTerminated,//关闭
    StateWait, // 等待
    StateTransition, // 切换
    StateUnknown
}THREAD_STATE;

// 线程处于等待的原因的枚举常量
typedef enum _KWAIT_REASON
{
    Executive,
    FreePage,
    PageIn,
    PoolAllocation,
    DelayExecution,
    Suspended,
    UserRequest,
    WrExecutive,
    WrFreePage,
    WrPageIn,
    WrPoolAllocation,
    WrDelayExecution,
    WrSuspended,
    WrUserRequest,
    WrEventPair,
    WrQueue,
    WrLpcReceive,
    WrLpcReply,
    WrVirtualMemory,
    WrPageOut,
    WrRendezvous,
    Spare2,
    Spare3,
    Spare4,
    Spare5,
    Spare6,
    WrKernel,
    MaximumWaitReason
}KWAIT_REASON;

typedef LONG NTSTATUS;
typedef LONG KPRIORITY;

typedef struct _CLIENT_ID
{
    DWORD UniqueProcess;
    DWORD UniqueThread;
} CLIENT_ID, *PCLIENT_ID;


typedef struct _VM_COUNTERS
{
    SIZE_T PeakVirtualSize;
    SIZE_T VirtualSize;
    ULONG  PageFaultCount;
    SIZE_T PeakWorkingSetSize;
    SIZE_T WorkingSetSize;
    SIZE_T QuotaPeakPagedPoolUsage;
    SIZE_T QuotaPagedPoolUsage;
    SIZE_T QuotaPeakNonPagedPoolUsage;
    SIZE_T QuotaNonPagedPoolUsage;
    SIZE_T PagefileUsage;
    SIZE_T PeakPagefileUsage;
} VM_COUNTERS;

// 线程信息结构体
typedef struct _SYSTEM_THREAD_INFORMATION
{
    LARGE_INTEGER KernelTime;
    LARGE_INTEGER UserTime;
    LARGE_INTEGER CreateTime;
    ULONG WaitTime;
    PVOID StartAddress;
    CLIENT_ID  ClientId;
    KPRIORITY Priority;
    KPRIORITY BasePriority;
    ULONG ContextSwitchCount;
    LONG  State;// 状态,是THREAD_STATE枚举类型中的一个值
    LONG WaitReason;//等待原因, KWAIT_REASON中的一个值
} SYSTEM_THREAD_INFORMATION, *PSYSTEM_THREAD_INFORMATION;


typedef struct _UNICODE_STRING
{
    USHORT Length;
    USHORT MaximumLength;
    PWSTR Buffer;
} UNICODE_STRING, *PUNICODE_STRING;


// 进程信息结构体
typedef struct _SYSTEM_PROCESS_INFORMATION
{
    ULONG            NextEntryDelta; // 指向下一个结构体的指针
    ULONG            ThreadCount; // 本进程的总线程数
    ULONG            Reserved1[6]; // 保留
    LARGE_INTEGER    CreateTime; // 进程的创建时间
    LARGE_INTEGER    UserTime; // 在用户层的使用时间
    LARGE_INTEGER    KernelTime; // 在内核层的使用时间
    UNICODE_STRING   ProcessName; // 进程名
    KPRIORITY        BasePriority; // 
    ULONG            ProcessId; // 进程ID
    ULONG            InheritedFromProcessId;
    ULONG            HandleCount; // 进程的句柄总数
    ULONG            Reserved2[2]; // 保留
    VM_COUNTERS      VmCounters;
    IO_COUNTERS      IoCounters;
    SYSTEM_THREAD_INFORMATION Threads[5]; // 子线程信息数组
}SYSTEM_PROCESS_INFORMATION, *PSYSTEM_PROCESS_INFORMATION;


// NtQuerySystemInformation 函数的原型
// 由于该没有导出,所以得自己定义函数的原型
typedef DWORD(WINAPI* PQUERYSYSTEM)(UINT, PVOID, DWORD, PDWORD);


DWORD __stdcall GetMainThread(DWORD dwProcessId)
{
    int nRet = 0;
    NTSTATUS Status = 0;
    DWORD dwThreadId = 0;

    PQUERYSYSTEM NtQuerySystemInformation = NULL;
    PSYSTEM_PROCESS_INFORMATION pInfo = { 0 };

    // 获取函数地址
    NtQuerySystemInformation = (PQUERYSYSTEM)GetProcAddress(LoadLibrary(L"ntdll.dll"), "NtQuerySystemInformation");


    DWORD dwSize = 0;
    // 获取信息所需的缓冲区大小
    Status = NtQuerySystemInformation(SystemProcessesAndThreadsInformation,//要获取的信息的类型
        NULL, // 用于接收信息的缓冲区
        0,
        &dwSize
        );
    // 申请缓冲区
    char* pBuff = new char[dwSize];
    pInfo = (PSYSTEM_PROCESS_INFORMATION)pBuff;
    if (pInfo == NULL)
        return -1;
    // 再次调用函数, 获取信息
    Status = NtQuerySystemInformation(SystemProcessesAndThreadsInformation,//要获取的信息的类型
        pInfo, // 用于接收信息的缓冲区
        dwSize,
        &dwSize
        );
    if (!NT_SUCCESS(Status)) {/*如果函数执行失败*/
        delete[] pInfo;
        return -1;
    }

    // 遍历结构体,找到对应的进程
    while (1)
    {
        // 判断是否还有下一个进程
        if (pInfo->NextEntryDelta == 0)
            break;

        // 判断是否找到了ID
        if (pInfo->ProcessId == dwProcessId)
        {
            dwThreadId = pInfo->Threads[0].ClientId.UniqueThread;
            break;
        }
        // 迭代到下一个节点
        pInfo = (PSYSTEM_PROCESS_INFORMATION)(((PUCHAR)pInfo) + pInfo->NextEntryDelta);
    }

    delete[] pBuff;
    return dwThreadId;
}


//打开调试权限
BOOL EnableDebugPrivilege(void)
{
    HANDLE hToken;
    TOKEN_PRIVILEGES token_p;
    if (OpenProcessToken(GetCurrentProcess(), TOKEN_ADJUST_PRIVILEGES | TOKEN_QUERY, &hToken)) {
        LookupPrivilegeValue(NULL, SE_DEBUG_NAME, &token_p.Privileges[0].Luid);//要求调试进程的权限 
        token_p.PrivilegeCount = 1;
        token_p.Privileges[0].Attributes = SE_PRIVILEGE_ENABLED;
        if (AdjustTokenPrivileges(hToken, FALSE, &token_p, sizeof(token_p), NULL, NULL))//调整权限
            return TRUE;
    }
    return FALSE;
}